import { Injectable } from '@nestjs/common';
import { getAgentWithConfig } from './utils/model';
import { createUIMessageStream } from 'ai';
import { createSSEData } from './utils/common';

@Injectable()
export class AiService {
  async test01() {
    const agent = await getAgentWithConfig({
      streaming: true,
    });
    function getAdditionalDataStream() {
      return new ReadableStream({
        start(controller) {
          controller.enqueue('这是来自附加流的数据块1。');
          controller.enqueue('这是来自附加流的数据块2。');
          controller.close();
        },
      });
    }
    const stream = createUIMessageStream({
      execute: async ({ writer }) => {
        const aiResponse = await agent.stream(
          {
            messages: [{ role: 'user', content: '写一个冒泡排序' }],
          },
          { streamMode: 'messages', configurable: { thread_id: 'xx' } },
        );
        for await (const chunk of aiResponse) {
          const [messageChunk] = chunk;
          if (!messageChunk.content) {
            // 合并附加数据流
            writer.merge(getAdditionalDataStream());
          }
          writer.write({
            type: 'data-appendMessage',
            data: {
              content: messageChunk.content,
            },
          });
        }
      },
    }).pipeThrough(
      new TransformStream({
        transform(chunk, controller) {
          controller.enqueue(createSSEData(chunk));
        },
      }),
    );
    return stream;
  }
}
